import cv2
import os
import numpy as np
import mediapipe as mp
import json
from datetime import datetime

# 创建数据集目录
os.makedirs("dynamic_dataset", exist_ok=True)
categories = [
    'right_angle', 'acute_angle', 'lightning',
    'triangle', 'H', 'R', 'W', 'PHI',
    'circle', 'wave'
]

# 初始化MediaPipe
mp_hands = mp.solutions.hands
hands = mp_hands.Hands(
    max_num_hands=1,
    min_detection_confidence=0.7,
    min_tracking_confidence=0.5
)

# 轨迹记录参数
MAX_POINTS = 100  # 最大轨迹点数
MIN_POINTS = 20  # 最小轨迹点数
TRAJECTORY = []  # 当前轨迹


def record_trajectory(frame, landmarks):
    """记录食指指尖的轨迹"""
    h, w, _ = frame.shape
    index_tip = landmarks.landmark[mp_hands.HandLandmark.INDEX_FINGER_TIP]
    x, y = int(index_tip.x * w), int(index_tip.y * h)

    # 只记录移动的点
    if not TRAJECTORY or (x, y) != TRAJECTORY[-1]:
        TRAJECTORY.append((x, y))

    # 绘制轨迹
    for i in range(1, len(TRAJECTORY)):
        cv2.line(frame, TRAJECTORY[i - 1], TRAJECTORY[i], (0, 255, 255), 2)


def save_trajectory(category):
    """保存轨迹到文件"""
    if len(TRAJECTORY) < MIN_POINTS:
        print(f"轨迹太短 ({len(TRAJECTORY)}点)，未保存")
        return False

    # 创建时间戳文件名
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
    filename = f"dynamic_dataset/{category}/{timestamp}.json"

    # 确保目录存在
    os.makedirs(os.path.dirname(filename), exist_ok=True)

    # 保存为JSON
    with open(filename, 'w') as f:
        json.dump({
            "category": category,
            "points": TRAJECTORY
        }, f)

    print(f"已保存 {category} 轨迹, {len(TRAJECTORY)}点")
    return True


def main():
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print("错误: 无法访问摄像头")
        return

    current_category = 0
    sample_count = {category: 0 for category in categories}
    recording = False

    # 创建键位映射
    key_map = {}
    key_labels = []

    # 前9个类别使用数字键1-9
    for i in range(min(9, len(categories))):
        key_map[str(i + 1)] = i
        key_labels.append(f"{i + 1}: {categories[i]}")

    # 第10个及之后的类别使用字母键a,b,c...
    for i in range(9, len(categories)):
        key_char = chr(ord('a') + i - 9)
        key_map[key_char] = i
        key_labels.append(f"{key_char}: {categories[i]}")

    print("动态手势数据集生成工具")
    print("按数字键选择类别：")
    for label in key_labels:
        print(label)
    print("按空格键开始/结束录制")
    print("按q退出")

    while cap.isOpened():
        success, frame = cap.read()
        if not success:
            print("错误: 无法读取摄像头帧")
            break

        frame = cv2.flip(frame, 1)
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        results = hands.process(rgb_frame)

        # 显示当前选择的类别
        category_name = categories[current_category]
        status = "录制中" if recording else "待机"
        cv2.putText(frame, f"类别: {category_name} ({sample_count[category_name]})",
                    (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
        cv2.putText(frame, f"状态: {status}, 点数: {len(TRAJECTORY)}",
                    (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2)
        cv2.putText(frame, "1-9/a-b: 选择类别, 空格: 开始/结束录制, q: 退出",
                    (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (200, 200, 255), 1)

        if results.multi_hand_landmarks:
            landmarks = results.multi_hand_landmarks[0]
            mp.solutions.drawing_utils.draw_landmarks(
                frame, landmarks, mp_hands.HAND_CONNECTIONS)

            # 录制时记录轨迹
            if recording:
                record_trajectory(frame, landmarks)

        # 显示轨迹
        for i in range(1, len(TRAJECTORY)):
            cv2.line(frame, TRAJECTORY[i - 1], TRAJECTORY[i], (0, 200, 255), 3)

        cv2.imshow('Dynamic Gesture Recorder', frame)

        key = cv2.waitKey(1)
        if key & 0xFF == ord('q'):
            break
        elif chr(key & 0xFF) in key_map:
            # 处理按键选择类别
            key_char = chr(key & 0xFF)
            current_category = key_map[key_char]
            print(f"切换到类别: {categories[current_category]}")
        elif key & 0xFF == 32:  # 空格键
            if recording:
                # 结束录制并保存
                if save_trajectory(category_name):
                    sample_count[category_name] += 1
                recording = False
                TRAJECTORY.clear()
            else:
                # 开始录制
                recording = True
                TRAJECTORY.clear()
                print(f"开始录制 {category_name}...")

    cap.release()
    cv2.destroyAllWindows()
    print("数据集生成完成！")


if __name__ == "__main__":
    main()